import torch
from einops import rearrange
from torch.nn import Linear, Parameter
from torch_geometric.nn import MessagePassing
from torch_geometric.utils import add_self_loops, degree


class SGCNConv(MessagePassing):
    def __init__(self, in_channels, out_channels):
        super().__init__(aggr='add')  # "Add" aggregation (Step 5).
        self.lin = Linear(in_channels, out_channels, bias=False)
        self.bias = Parameter(torch.empty(out_channels))

        self.reset_parameters()

    def reset_parameters(self):
        self.lin.reset_parameters()
        self.bias.data.zero_()

    def forward(self, x, edge_index):
        # x has shape [N, in_channels]
        # edge_index has shape [2, E]

        # Step 1: Add self-loops to the adjacency matrix.
        edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))

        # Step 2: Linearly transform node feature matrix.
        x = self.lin(x)

        # Step 3: Compute normalization.
        row, col = edge_index
        deg = degree(col, x.size(0), dtype=x.dtype)
        deg_inv_sqrt = deg.pow(-0.5)
        # deg_sqrt = deg ** -0.5
        # f = torch.equal(deg_inv_sqrt,deg_sqrt)
        deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
        norm = deg_inv_sqrt[row] * deg_inv_sqrt[col]

        # Step 4-5: Start propagating messages.
        t = norm.view(-1, 1)
        t1 = rearrange(norm, "i -> i ()")  # [E+N,1]
        # out1 = torch.einsum("i j,n m->n m", t1, x)

        out = self.propagate(edge_index, x=x, norm=norm)

        # f = torch.equal(out1, out)

        # Step 6: Apply a final bias vector.
        out += self.bias

        return out

    def message(self, x_j, norm):
        # x_j has shape [E, out_channels]

        # Step 4: Normalize node features.
        return norm.view(-1, 1) * x_j


if __name__ == "__main__":
    h = torch.randn(4, 8)  # 4个节点，每个节点维度为8的特征
    edge_index = torch.tensor([[0, 1],
                               [1, 0]])
    # 对边矩阵(2,E) E为边的数量，第一行为边的起点，第二行为边的目标节点

    gcn = GCNConv(in_channels=8, out_channels=16)
    h1 = gcn(h, edge_index)

    pass
